import { ChatOllama } from "@langchain/ollama";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { StringOutputParser } from "@langchain/core/output_parsers";
import { RunnableLambda, RunnableSequence } from "@langchain/core/runnables";

// 1. 模型
const model = new ChatOllama({
  model: "llama3",
});

// 2. 解析器
const parser = new StringOutputParser();

// 3. 三条“领域链” —— 物理 / 数学 / 其他
const physicsChain = ChatPromptTemplate.fromTemplate(
  `你是一位物理学家，擅长使用中文回答物理相关的问题，当你不知道问题的答案时，你就回答不知道。
具体问题如下：
{input}`
)
  .pipe(model)
  .pipe(parser);

const mathChain = ChatPromptTemplate.fromTemplate(
  `你是一个数学家，擅长使用中文回答数学相关的问题，当你不知道问题的答案时，你就回答不知道。
具体问题如下：
{input}`
)
  .pipe(model)
  .pipe(parser);

const otherChain = ChatPromptTemplate.fromTemplate(
  `你是一个AI助手，你会使用中文回答以下问题。
具体问题如下：
{input}`
)
  .pipe(model)
  .pipe(parser);

// 4. 分类链：输出“数学”/“物理”/“其它”
const classifyChain = ChatPromptTemplate.fromMessages([
  [
    "system",
    `
你是一个仅输出标签的分类器。
任务：将用户问题分类为以下三类之一：数学、物理、其它。
规则：
- 只能输出且必须输出其中一个：数学 或 物理 或 其它
- 不要输出任何解释、理由、引号、标点、前后缀或额外空行
- 若不确定，输出：其它
    `.trim(),
  ],
  [
    "human",
    `
问题：
{input}

现在输出分类（仅一个词）：`.trim(),
  ],
])
  .pipe(model)
  .pipe(parser); // 分类链

// 路由器
const router = RunnableLambda.from(async ({ input }) => {
  // input : 入参 - “什么是经典力学？”、“对 y = x 求导的结果是多少”
  const topic = await classifyChain.invoke({
    input,
  });
  console.log("当前的分类结果：", topic); // “当前的分类结果： 物理”

  // 根据分类结果，选择一条链条
  const chosenChain = topic.includes("数学")
    ? mathChain
    : topic.includes("物理")
    ? physicsChain
    : otherChain;

  return await chosenChain.invoke({
    input,
  });
});

// 5. 组合为完整链
const finalChain = RunnableSequence.from([
  RunnableLambda.from((input) => input),
  router,
]);

// 测试
const run = async () => {
  console.log(await finalChain.invoke({ input: "什么是经典力学？" }));
  console.log(
    await finalChain.invoke({ input: "对 y = x 求导的结果是多少？" })
  );
  console.log(await finalChain.invoke({ input: "你好，你是谁？" }));
};
run();
